from deoldify import device
from deoldify.device_id import DeviceId
from deoldify.visualize import get_image_colorizer
import gradio as gr
import warnings

# 忽略警告信息
warnings.filterwarnings("ignore", category=UserWarning, message=".*?Your .*? set is empty.*?")

# 设置设备为 GPU，如果可用
device.set(device=DeviceId.GPU0)

# 加载预训练的 deoldify 模型
colorizer = get_image_colorizer(artistic=True)

def colorize_image(input_img, render_factor, watermarked):
    # 上色处理
    colorized_img = colorizer.get_transformed_image(input_img, render_factor=render_factor, watermarked=watermarked)
    return colorized_img

# 创建 Gradio 应用
iface = gr.Interface(
    fn=colorize_image,                                        # 处理函数
    inputs=[
        gr.Image(type="pil", label="黑白图像"),                 # 黑白图像
        gr.Slider(7, 40, step=1, value=26, label="渲染系数 "),  # 渲染因子
    ],
    outputs=gr.Image(type="pil", label="彩色图像"),             # 彩色图像
    title="图像上色器",                                         # 标题
    description="上传一张黑白照片以通过Deoldify为其上色。",         # 应用描述
    theme="default"                                           # Gradio主题
)

# 启动应用
iface.launch()